import asyncio
import time
import numpy as np
import pandas as pd
import sqlite3
from tqsdk import TqApi, TqAuth, TqKq
from MyTT import EMA

class TradingBot:
    def __init__(self):
        self.api = None
        self.running = False
        self.active_symbols = set()
        self.tasks = []
        
    async def initialize(self):
        """初始化API连接"""
        try:
            self.api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba999"))
            await self.api.wait_update()
            return True
        except Exception as e:
            print(f"API初始化失败: {e}")
            return False
    
    async def get_monitored_symbols(self):
        """从数据库获取监控品种"""
        conn = sqlite3.connect('trade_monitor.db')
        cursor = conn.cursor()
        cursor.execute("SELECT symbol FROM products WHERE is_monitored=1")
        symbols = [row[0] for row in cursor.fetchall()]
        conn.close()
        return symbols
    
    async def cleanup(self):
        """改进的资源清理方法"""
        self.running = False
        
        # 取消所有任务
        for task in self.tasks:
            task.cancel()
            try:
                await task
            except (asyncio.CancelledError, Exception):
                pass
                
        # 等待所有活动品种处理完成
        while len(self.active_symbols) > 0:
            await asyncio.sleep(0.1)
            
        # 关闭API连接
        if self.api:
            try:
                # 创建新的事件循环用于关闭API
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                await self.api.close()
                loop.close()
            except Exception:
                pass
                
        self.tasks = []

    async def main_listener(self):
        """主监听循环"""
        async with self.api.register_update_notify() as update_chan:
            async for _ in update_chan:
                if not self.running:
                    break

    async def run(self):
        """改进的主运行逻辑"""
        if not await self.initialize():
            return
            
        self.running = True
        symbols = await self.get_monitored_symbols()
        
        try:
            # 使用asyncio.create_task创建任务
            self.tasks = [asyncio.create_task(self.process_symbol(symbol))
                         for symbol in symbols]
            
            # 添加主监听任务
            main_task = asyncio.create_task(self.main_listener())
            self.tasks.append(main_task)
            
            await asyncio.gather(*self.tasks)
        except Exception as e:
            print(f"策略运行出错: {e}")
        finally:
            await self.cleanup()

    # 确保以下方法都在类定义内部，缩进正确
    async def process_symbol(self, symbol):
        """处理单个品种"""
        self.active_symbols.add(symbol)
        try:
            quote = self.api.get_quote(symbol)
            klines = self.api.get_kline_serial(symbol, 60, 500)  # 1分钟K线
            
            async with self.api.register_update_notify() as update_chan:
                async for _ in update_chan:
                    if not self.running:
                        break
                        
                    if self.api.is_changing(quote):
                        await self.check_stoploss(symbol, quote, klines)
                        
        except Exception as e:
            print(f"处理品种 {symbol} 时出错: {e}")
        finally:
            self.active_symbols.remove(symbol)

async def check_stoploss(self, symbol, quote, klines):
    """检查并执行止损逻辑"""
    print(f"{time.strftime('%X')} {symbol} 行情更新，当前价: {quote.last_price}")
    
    # 计算技术指标
    q = (3*klines.close + klines.low + klines.open + klines.high) / 6
    ma_values = self.calculate_ma(q)
    short_ma = ma_values['ma']
    ema_7 = ma_values['ema']
    
    position = self.api.get_position(symbol)
    
    # 平空仓逻辑
    if position.pos_short > 0 and short_ma[-2] > ema_7[-2]:
        await self.api.insert_order(
            symbol=symbol,
            direction="BUY",
            offset="CLOSE",
            volume=position.pos_short
        )
        print(f"{time.strftime('%X')} 平空仓: {symbol} 手数: {position.pos_short}")
    
    # 平多仓逻辑
    if position.pos_long > 0 and short_ma[-2] < ema_7[-2]:
        await self.api.insert_order(
            symbol=symbol,
            direction="SELL",
            offset="CLOSE",
            volume=position.pos_long
        )
        print(f"{time.strftime('%X')} 平多仓: {symbol} 手数: {position.pos_long}")

def calculate_ma(self, q):
    """计算移动平均指标"""
    weights = np.arange(26, 0, -1)
    ma = []
    q_array = np.array(q)
    for i in range(len(q_array)):
        if i < 25:
            ma.append(np.nan)
            continue
        window = q_array[i-25:i+1]
        ma.append(np.dot(window, weights) / 351)
    
    ema_7 = EMA(pd.Series(ma), 7)
    return {'ma': ma, 'ema': ema_7.tolist()}

if __name__ == "__main__":
    print("****************************************")
    print("****** 天勤量化移动止损程序 v3.0 *********")
    print("****************************************")
    
    bot = TradingBot()
    try:
        asyncio.run(bot.run())
    except KeyboardInterrupt:
        print("\n程序已手动终止")
